package syncx

import (
	"sync"
	"time"
)

type WaitGroup struct {
	sync.WaitGroup
}

// WaitWithTimeout blocks current goroutine until wg.Done() or timeout.
func (wg *WaitGroup) WaitWithTimeout(d time.Duration) bool {
	done := make(chan struct{}, 1) // done when wg.Done()
	go func() {
		defer close(done)
		wg.Wait()
		done <- struct{}{}
	}()

	select {
	case <-done:
		return true
	case <-time.After(d):
		return false
	}
}
